package com.blackholemax.algorithm.search;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

public class TrieDfs {
    // 定义字典树节点类 TrieNode，用于存储字典树的每个节点
    private static class TrieNode {
        // 标志当前节点是否是一个单词的结尾
        boolean isEnd;
        // 使用 TreeMap 存储子节点，TreeMap 自动按字典序排序
        TreeMap<Character, TrieNode> children;

        // 构造函数，初始化节点
        TrieNode() {
            isEnd = false; // 初始化时，默认不是单词结尾
            children = new TreeMap<>(); // 初始化子节点为空的 TreeMap
        }
    }

    // 定义字典树类 Trie，用于实现插入和前缀搜索功能
    private static class Trie {
        // 定义根节点，作为字典树的入口
        private TrieNode root;

        // 构造函数，初始化字典树
        Trie() {
            root = new TrieNode(); // 创建根节点
        }

        // 插入单词的方法
        public void insert(String word) {
            TrieNode node = root; // 从根节点开始
            for (char c : word.toCharArray()) { // 遍历单词的每个字符
                if (!node.children.containsKey(c)) { // 如果当前节点没有该字符对应的子节点
                    node.children.put(c, new TrieNode()); // 创建新的子节点
                }
                node = node.children.get(c); // 移动到子节点
            }
            node.isEnd = true; // 标记单词的最后一个字符为单词结尾
        }

        // 前缀搜索方法，返回所有以指定前缀开头的单词
        public List<String> search(String prefix) {
            List<String> result = new ArrayList<>(); // 用于存储搜索结果
            TrieNode node = root; // 从根节点开始
            for (char c : prefix.toCharArray()) { // 遍历前缀的每个字符
                if (!node.children.containsKey(c)) { // 如果前缀不存在于字典树中
                    return result; // 返回空列表
                }
                node = node.children.get(c); // 移动到子节点
            }
            // 调用 dfs 方法，从当前节点开始深度优先遍历，收集所有匹配的单词
            dfs(node, prefix, result);
            return result; // 返回搜索结果
        }

        // 深度优先遍历方法，用于收集所有以指定前缀开头的单词
        private void dfs(TrieNode node, String path, List<String> result) {
            if (node.isEnd) { // 如果当前节点是一个单词的结尾
                result.add(path); // 将当前路径（即单词）加入结果列表
            }
            // 遍历当前节点的所有子节点
            for (Map.Entry<Character, TrieNode> entry : node.children.entrySet()) {
                // 递归调用 dfs 方法，继续遍历子节点
                dfs(entry.getValue(), path + entry.getKey(), result);
            }
        }
    }

    // 主方法，程序入口
    public static void main(String[] args) {
        Trie trie = new Trie(); // 创建字典树实例
        trie.insert("apple"); // 插入单词 "apple"
        trie.insert("app"); // 插入单词 "app"
        trie.insert("banana"); // 插入单词 "banana"
        trie.insert("bat"); // 插入单词 "bat"

        // 前缀搜索 "ap"，打印结果
        System.out.println(trie.search("ap")); // 输出: [app, apple]
        trie.insert("apply"); // 插入新单词 "apply"
        // 再次前缀搜索 "ap"，打印结果
        System.out.println(trie.search("ap")); // 输出: [app, apple, apply]
    }
}

